"""
Methods for performing Gaussian Quadrature

Functions
---------
quad            Clenshaw-Curtis or Gaussian quadrature
golub_welsch    General tool for generating general quadrature
                nodes, weights
stieltjes       Tool for performing discretized and analytical
                Stieltjes procedure.
"""

import numpy as np
from scipy.misc import comb
import bertran as ber
import poly as po
from utils import combine, lazy_eval
import dist as di

__version__ = "2.0"

def golub_welsch(order, dist, acc=100):
    """
Golub-Welsch algorithm for creating quadrature nodes and weights

Parameters
----------
order : int
    Quadrature order
dist : Dist
    Distribution nodes and weights are found for with
    `dim=len(dist)`
acc : int
    Accuracy used in discretized Stieltjes procedure.
    Will be increased by one for each itteration.

Returns
-------
x : numpy.array
    Optimal collocation nodes with `x.shape=(dim, order+1)`
w : numpy.array
    Optimal collocation weights with `w.shape=(order+1,)`

Examples
--------
>>> Z = pc.Normal()
>>> x, w = pc.golub_welsch(3, Z)
>>> print x
[[-2.33441422 -0.74196378  0.74196378  2.33441422]]
>>> print w
[ 0.04587585  0.45412415  0.45412415  0.04587585]

Multivariate
>>> Z = pc.J(pc.Uniform(), pc.Uniform())
>>> x, w = pc. golub_welsch(1, Z)
>>> print x
[[ 0.21132487  0.21132487  0.78867513  0.78867513]
 [ 0.21132487  0.78867513  0.21132487  0.78867513]]
>>> print w
[ 0.25  0.25  0.25  0.25]
    """

    P,g,a,b = stieltjes(dist, order, acc, True)

    X,W = [], []
    dim = len(dist)
    for d in xrange(dim):
        J = np.diag(np.sqrt(b[d,1:]), k=-1) + np.diag(a[d]) + \
                np.diag(np.sqrt(b[d,1:]), k=1)
        vals, vecs = np.linalg.eig(J)

        x, w = vals.real, vecs[0,:]**2
        indices = np.argsort(x)
        x, w = x[indices], w[indices]
        x[np.abs(x)<1e-14] = 0

        X.append(x)
        W.append(w)

    x = combine(X)
    w = combine(W)

    w = np.prod(w, -1)
    return x.T, w.T


def stieltjes(dist, order, normed=False, retall=False, acc=100,
        **kws):
    """
Discretized Stieltjes method

Parameters
----------
dist : Dist
    Distribution defining the space to create weights for.
order : int
    Number of iterations to create.
acc : int
    The number of Clenshaw-Curtis nodes to use at each step.
retall : bool
    If included, more values are returned

Returns
-------
orth[, norms, A, B]

orth : list
    List of polynomials created from the method with
    `len(orth)==order+1`.
    If `len(dist)>1`, then each polynomials are multivariate.
norms : np.ndarray
    The norms of the polynomials with `norms.shape=(dim,order+1)`
    where `dim` are the number of dimensions in dist.
A,B : np.ndarray
    The three term coefficients. Both have `shape=(dim,order+1)`.

Examples
--------
>>> dist = pc.Uniform()
>>> orth, norms, A, B = pc.stieltjes(dist, 2, retall=True)
>>> print orth[2]
[q0^2-q0+0.166666666667]
>>> print norms
[[ 1.          0.08333333  0.00555556]]
    """

    try:
        dim = len(dist)
        K = np.arange(order+1).repeat(dim).reshape(order+1, dim).T
        A,B = dist.ttr(K)
        B[:,0] = 1.

        x = po.variable(dim)
        if normed:
            orth = [x**0*np.ones(dim), (x-A[:,0])/np.sqrt(B[:,1])]
            for n in xrange(1,order):
                orth.append((orth[-1]*(x-A[:,n]) - orth[-2]*np.sqrt(B[:,n]))/np.sqrt(B[:,n+1]))
            norms = np.ones(B.shape)
        else:
            orth = [x-x, x**0*np.ones(dim)]
            for n in xrange(order):
                orth.append(orth[-1]*(x-A[:,n]) - orth[-2]*B[:,n])
            orth = orth[1:]
            norms = np.cumprod(B, 1)

    except NotImplementedError:

        bnd = dist.range()
        q,w = quadgen(acc, bnd[0], bnd[1])
        w = w*dist.pdf(q)

        dim = len(dist)
        x = po.variable(dim)
        orth = [x*0, x**0]

        inner = np.sum(q*w, -1)
        norms = [np.ones(dim), np.ones(dim)]
        A,B = [],[]

        for n in xrange(order):

            A.append(inner/norms[-1])
            B.append(norms[-1]/norms[-2])
            orth.append((x-A[-1])*orth[-1] - orth[-2]*B[-1])

            y = orth[-1](*q)**2*w
            inner = np.sum(q*y, -1)
            norms.append(np.sum(y, -1))

        A, B = np.array(A).T, np.array(B).T
        norms = np.array(norms[1:]).T
        orth = orth[1:]

    if retall:
        return orth, norms, A, B
    return orth


def _clenshaw_curtis(N, composit=1):

    if not N:
        return np.array([0.]), np.array([2.])

    N = 2*int(N)
    x = np.cos(np.arange(N/2+1)*np.pi/N)
    x = np.concatenate((-x[:-1*(N%2==0)], x[::-1]))
    x[np.abs(x)<1e-14] = 0

    k,n = np.meshgrid(*[np.arange(N/2+1)]*2)
    D = 2./N*np.cos(2*n*k*np.pi/N)
    D[:,0] *= .5
    D[:,-1] *= .5

    d = 2./(1-np.arange(0,N+1,2)**2)
    d[0] *= .5
    d[-1] *= .5

    w = np.dot(D.T, d)
    w = np.concatenate((w, w[-1-1*(N%2==0)::-1]))
    w[N/2] *= 2

    x = .5*x+.5
    w *= .5

    M = len(x)


    composit = list(set(composit))
    composit.sort()

    if composit[0]!=0:
        composit = [0]+composit
    if composit[-1]!=1:
        composit = composit+[1]

    X = np.zeros((M-1)*(len(composit)-1)+1)
    W = np.zeros((M-1)*(len(composit)-1)+1)
    for d in xrange(len(composit)-1):
        X[d*M-d:(d+1)*M-d] = \
                x*(composit[d+1]-composit[d]) + composit[d]
        W[d*M-d:(d+1)*M-d] += w*(composit[d+1]-composit[d])

    return X, W


def clenshaw_curtis(N, lo=0, up=1, growth=0, composit=1):
    """
Generate the quadrature nodes and weights in Clenshaw-Curtis
quadrature
    """
    N,lo,up = [np.array(_).flatten() for _ in [N,lo,up]]
    dim = max(lo.size, up.size, N.size)
    N,lo,up = [np.ones(dim)*_ for _ in [N,lo,up]]
    N = np.array(N, dtype=int)

    if isinstance(composit, int):
        composit = [np.linspace(0,1,composit+1)]*dim
    else:
        composit = np.array(composit)
        if not composit.shape:
            composit = composit.flatten()
        if len(composit.shape)==1:
            composit = np.array([composit]).T
        composit = ((composit.T-lo)/(up-lo)).T

    if growth==0:
        q = [_clenshaw_curtis((N[i]+1)/2, composit[i]) for i in xrange(dim)]
    elif growth==1:
        q = [_clenshaw_curtis(N[i], composit[i]) for i in xrange(dim)]
    elif growth==2:
        q = [_clenshaw_curtis(2**(N[i]-1)-.5*(N[i]==0), composit[i]) \
                for i in xrange(dim)]

    x = [_[0] for _ in q]
    w = [_[1] for _ in q]

    x = combine(x)
    w = combine(w)

    x = (up-lo)*x + lo
    w = np.prod(w*(up-lo), 1)

    return x.T, w


def _gauss_legendre(N, composit=1):
    a = np.ones(N+1)*0.5
    b = np.arange(N+1)**2
    b = b/(16*b-4.)

    J = np.diag(np.sqrt(b[1:]), k=-1) + np.diag(a) + \
            np.diag(np.sqrt(b[1:]), k=1)
    vals, vecs = np.linalg.eig(J)

    x, w = vals.real, vecs[0,:]**2
    indices = np.argsort(x)
    x, w = x[indices], w[indices]

    X = np.zeros((N+1)*composit)
    W = np.zeros((N+1)*composit)

    for d in xrange(composit):
        X[d*(N+1):(d+1)*(N+1)] = (x+d)/composit
        W[d*(N+1):(d+1)*(N+1)] = w/composit

    return X, W

def gauss_legendre(N, lo=0, up=1, composit=1):
    """
Generate the quadrature nodes and weights in Gauss-Legendre
quadrature
    """

    N,lo,up = [np.array(_).flatten() for _ in [N,lo,up]]
    dim = max(lo.size, up.size, N.size)
    N,lo,up = [np.ones(dim)*_ for _ in [N,lo,up]]
    N = np.array(N, dtype=int)

    q = [_gauss_legendre(N[i], composit) for i in xrange(dim)]
    x = np.array([_[0] for _ in q])
    w = np.array([_[1] for _ in q])

    x = combine(x)
    w = combine(w)

    x = (up-lo)*x + lo
    w = np.prod(w*(up-lo), 1)

    return x.T, w


def sparse_grid(func, order, dim, skew=None):

    X, W = [], []
    bindex = ber.bindex(order-dim+1, order, dim)

    if skew is None:
        skew = np.zeros(dim, dtype=int)
    else:
        skew = np.array(skew, dtype=int)
        assert len(skew)==dim

    for i in xrange(ber.terms(order, dim)-ber.terms(order-dim, dim)):

        I = bindex[i]
        x,w = func(skew+I)
        w *= (-1)**(order-sum(I))*comb(dim-1,order-sum(I))
        X.append(x)
        W.append(w)

    X = np.concatenate(X, 1)
    W = np.concatenate(W, 0)

    order = np.lexsort(X)
    X = X.T[order].T
    W = W[order]

    # identify non-unique terms
    diff = np.diff(X.T, axis=0)
    ui = np.ones(len(X.T), bool)
    ui[1:] = (diff!=0).any(axis=1)

    # merge duplicate nodes
    N = len(W)
    i = 1
    while i<N:
        while i<N and ui[i]: i+=1
        j = i+1
        while j<N and not ui[j]: j+=1
        if j-i>1:
            W[i-1] = np.sum(W[i-1:j])
        i = j+1

    X = X[:,ui]
    W = W[ui]

    return X, W

def momgen(func, order, domain, acc=100, sparse=False, rule="C",
        composit=1, **kws):
    X,W = quadgen(order, domain, acc, sparse, rule, composit, **kws)
    Y = np.array(func(X))

    def _mom(k):
        out = np.sum(np.prod(Y.T**k, -1)*W, 0)
        return out
    _mom = lazy_eval(_mom, tuple)

    def mom(K, **kws):
        out = np.array([_mom(k) for k in K.T])
        return out
    return mom

def quadgen(order, domain, acc=100, sparse=False, rule="C",
        composit=1, **kws):
    """
Numerical quadrature node and weight generator 

Parameters
----------
order : int
    The order of the quadrature.
domain : array_like, Dist
    If array is provided domain is the lower and upper bounds
    (lo,up). Invalid if gaussian is set.
    If Dist is provided, bounds and nodes are adapted to the
    distribution. This includes weighting the nodes in
    Clenshaw-Curtis quadrature.
acc : int
    If gaussian is set, but the Dist provieded in domain does not
    provide an analytical TTR, ac sets the approximation order for
    the descitized Stieltje's method.
sparse : bool
    If True used Smolyak's sparse grid instead of normal tensor
    product grid.
rule : str
    Quadrature rule

    Key     Description
    "G"     Optiomal Gaussian quadrature from Golub-Welsch
            Slow for high order and composit is ignored.
    "E"     Gauss-Legendre quadrature
    "C"     Clenshaw-Curtis quadrature. Exponential growth rule is
            used when sparse is True to make the rule nested.

    If other is provided, Monte Carlo integration is assumed and
    arguemnt is passed to samplegen with order and domain.

composit : int, optional
    If provided, composit quadrature will be used. Value determines
    the number of domains along an axis.
    Ignored in the case if gaussian=True.
**kws : optional
    Extra keyword passed to samplegen.

See also
--------
samplegen   Sample generator
    """

    rule = rule.upper()
    if rule=="G":

        assert isinstance(domain, di.Dist)

        if sparse:
            func = lambda m: golub_welsch(m, domain, acc)
            x, w = sparse_grid(func, order, len(domain))
        else:
            x, w = golub_welsch(order, domain, acc)


    elif rule in "EC":

        isdist = not isinstance(domain, (tuple, list, np.ndarray))

        if isdist:
            lo,up = domain.range()
            dim = len(domain)
        else:
            lo,up = np.array(domain)
            dim = lo.size

        if sparse:
            if rule=="C":
                func = lambda m: clenshaw_curtis(m, lo, up,
                        growth=2, composit=composit)
            elif rule=="E":
                func = lambda m: gauss_legendre(m, lo, up,
                        composit)
            x, w = sparse_grid(func, order, dim)
        else:
            if rule=="C":
                x, w = clenshaw_curtis(order, lo, up, growth=0,
                        composit=composit)
            elif rule=="E":
                x, w = gauss_legendre(order, lo, up, composit)

        if isdist:
            w *= domain.pdf(x).flatten()
        valid = np.abs(w)>1e-10/len(w)
        x, w = x[:,valid], w[valid]
        w /= np.sum(w)

    else:

        x = di.samplegen(order, domain, rule, **kws)
        w = np.ones(x.shape[-1])/x.shape[-1]

    return x, w

def quad(func, order, domain, acc=100, sparse=False,
        rule="C", composit=1, retall=False, **kws):
    """
Numerical integration using Quadrature of problems on the form
::int dist.pdf(x) func(x) dx::

Weights and nodes are either calculated using Golub Welsch or
Clenshaw-Curtis methods.
In the former case the recurrence coefficients are calculated
analytical if available or else from discretized Stieltjes
procedure using Clenshaw-Curtis quadrature.

Parameters
----------
func : callable(x)
    function to take quadrature over
    x : np.ndarray
        where x.shape=M.shape return shape Q must be consistent
domain : Dist
    Distribution to take density function used as weighting
    function, where `dim=len(dist)`
order : int, array_like
    Quadrature order
    If array provided, all elements must be non-negative int
    If int provided, it is converted to a length dim array.
acc : int
    Accuracy of discretized Stieltjes procedure.
rule : str
    Quadrature rule

    Key     Description
    "G"     Optiomal Gaussian quadrature from Golub-Welsch
            Slow for high order and composit is ignored.
    "L"     Gauss-Legendre quadrature
    "C"     Clenshaw-Curtis quadrature. Exponential growth rule is
            used when sparse is True to make the rule nested.

    If other is provided, Monte Carlo integration is assumed and
    passed to samplegen with order and domain.

composit : int, optional
    If provided, composit quadrature will be used. Value determines
    the number of domains along an axis.
    Ignored in the case if gaussian=True.
retall : bool
    If True return nodes, weights and evaluvation in addition to
    answer.
**kws : optional
    Extra keyword passed to samplegen.

Returns
-------
I[, x, w, y]

I : np.ndarray
    Integration estimate
x : np.ndarray
    Quadrature node with `x.shape=(K,D)` where `D=len(dist)`.
w : np.ndarray
    Quadrature weights with `x.shape=(K,D)` where `D=len(dist)`.
y : np.ndarray
    Function evaluation of func in x

Examples
--------
>>> dist = pc.Gamma()
>>> func = lambda x: x**3-1
>>> q, x, w, y = pc.quad(func, 1, dist, retall=True, rule="G")
>>> print q
[ 5.]
>>> print x
[[ 0.58578644  3.41421356]]
    """
    x,w = quadgen(order, domain, acc=acc, sparse=sparse, rule=rule,
            composit=composit)
    y = np.array(map(func, x.T))
    q = np.sum((y.T*w).T, 0)
    if retall:
        return q, x, w, y
    return q


def dep_golub_welsch(dist, order, acc=100):

    if not dist.dependent:
        return golub_welsch(dist, order, acc)
    raise NotImplementedError


#  def _quad_dependent(func, order, dist, acc=40, orth=None,
#      args=(), kws={}, veceval=False, retall=False):
#  
#      dim = len(dist)
#      if isinstance(order, (int, long, float)):
#          order = [order]*dim
#      indices = di.sort(dist)
#      grid = np.mgrid[[slice(0,order[i]+1,1) for i in indices]]
#      X = np.empty([dim,]+[order[i]+1 for i in indices])
#      W = np.ones([order[i]+1 for i in indices])
#  
#      def _dep_quad(I, order, dist, X, W, grid, vals):
#  
#          dim = len(dist)
#          i = len(I)
#          j = indices[i]
#  
#          Z = di.Subset(dist, j, vals)
#          x,w = Golub_Welsch(order[j], Z, acc)
#          W[I] *= w[grid[(i,)+I]]
#          X[(i,)+I] = x[grid[(i,)+I]]
#  
#          if i==dim-1:
#              return X, W
#  
#          for k in xrange(order[j]+1):
#              vals[j] = x[k]
#              X, W = _dep_quad(I+(k,), order, dist, X, W, grid, vals)
#  
#          return X,W
#  
#      X,W = _dep_quad((), order, dist, X, W, grid, [0.]*dim)
#      X = X[indices]
#      X = X.reshape(dim, X.size/dim)
#      W = W.flatten()
#  
#      if veceval:
#          Y = np.array(func(X, *args, **kws)).T
#      else:
#          Y = np.array([func(_, *args, **kws) for _ in X.T]).T
#  
#      if not (orth is None):
#          Q = orth(*X)
#  
#          shape = Y.shape[:-1] + Q.shape
#          Y = Y.reshape(np.prod(Y.shape[:-1]), Y.shape[-1],)
#          Q = Q.reshape(np.prod(Q.shape[:-1]), Q.shape[-1],)
#  
#          t1,t2 = np.mgrid[:len(Y), :len(Q)]
#          Y = (Y[t1]*Q[t2]).reshape(shape)
#  
#      out = np.sum(Y*W, -1).T
#  
#      if retall:
#          return out, X.T
#      return out




if __name__=="__main__":
    import __init__ as pc
    x,y = pc.variable(2)
    import doctest
    doctest.testmod()

